/**
 * @file coroutine.cpp
 * @author your name (you@domain.com)
 * @brief 
 * @version 0.1
 * @date 2024-06-18
 * 
 * @copyright Copyright (c) 2024
 * 
 */
#include "coroutine.hpp"

namespace libco {

/// 全局静态变量，用于生成协程id
static std::atomic<uint64_t> s_coroutine_id{1};
/// 全局静态变量，用于统计当前的协程数
static std::atomic<uint64_t> s_coroutine_count{0};
//协程栈大小,默认128k
static uint32_t default_stack_size = 128 * 1024;

/// 线程局部变量，当前线程正在运行的协程
static thread_local Coroutine *t_coroutine = nullptr;
/// 线程局部变量，当前线程的主协程，切换到这个协程，就相当于切换到了主线程中运行，智能指针形式
static thread_local Coroutine::ptr t_thread_coroutine = nullptr;


/**
 * @brief malloc栈内存分配器
 */
class StackAllocator {
public:
    static void *Alloc(size_t size) { return malloc(size); }
    static void Dealloc(void *vp, size_t size) { return free(vp); }
};

uint64_t Coroutine::GetCoroutineId() {
    if (t_coroutine) {
        return t_coroutine->getId();
    }
    return 0;
}

std::string Coroutine::GetCoroutineRole() {
    if (t_coroutine) {
        return t_coroutine->m_is_scheduler ? "scheduler" : "worker";
    }
    return "nocoroutine";
}

uint64_t Coroutine::TotalCoroutines()
{
    return s_coroutine_count;
}

Coroutine::Coroutine() {
    SetThis(this);
    m_state = RUNNING;
    m_is_scheduler = true;

    if (getcontext(&m_ctx)) {
        LIBCO_ASSERT2(false, "getcontext");
    }

    ++s_coroutine_count;
    m_id = s_coroutine_id++; // 协程id从0开始，用完加1

    LIBCO_LOG << "scheduler coroutine id = " << m_id << " begin";
}

void Coroutine::SetThis(Coroutine *co) { 
    t_coroutine = co; 
}

/**
 * 获取当前协程，同时充当初始化当前线程主协程的作用，这个函数在使用协程之前要调用一下
 */
Coroutine::ptr Coroutine::GetThis() {
    if (t_coroutine) {
        return t_coroutine->shared_from_this();
    }

    Coroutine::ptr main_coroutine(new Coroutine);
    LIBCO_ASSERT(t_coroutine == main_coroutine.get());
    t_thread_coroutine = main_coroutine;
    return t_coroutine->shared_from_this();
}

/**
 * 带参数的构造函数用于创建其他协程，需要分配栈
 */
Coroutine::Coroutine(std::function<void()> cb, size_t stacksize)
    : m_id(s_coroutine_id++)
    , m_cb(cb) {
    ++s_coroutine_count;
    m_stacksize = stacksize ? stacksize : default_stack_size;
    m_stack     = StackAllocator::Alloc(m_stacksize);
    m_is_scheduler = false;

    if (getcontext(&m_ctx)) {
        LIBCO_ASSERT2(false, "getcontext");
    }

    m_ctx.uc_link          = nullptr;
    m_ctx.uc_stack.ss_sp   = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;

    makecontext(&m_ctx, &Coroutine::MainFunc, 0);

    LIBCO_LOG_DEBUG << "create worker coroutine id = " << m_id;
}

/**
 * 线程的主协程析构时需要特殊处理，因为主协程没有分配栈和cb
 */
Coroutine::~Coroutine() {
    --s_coroutine_count;
    if (m_stack) {
        // 有栈，说明是子协程，需要确保子协程一定是结束状态

        LIBCO_ASSERT(m_state == TERM);
        StackAllocator::Dealloc(m_stack, m_stacksize);
        LIBCO_LOG_DEBUG << "destroy worker coroutine id = " << m_id;
    } else {
        // 没有栈，说明是线程的主协程
        LIBCO_ASSERT(!m_cb);              // 主协程没有cb
        LIBCO_ASSERT(m_state == RUNNING); // 主协程一定是执行状态
        LIBCO_ASSERT(t_coroutine == this); // 当前运行协程一定是调度协程

        LIBCO_LOG_DEBUG << "scheduler coroutine id = " << m_id << " ends";
        SetThis(nullptr);
        t_thread_coroutine = nullptr;
    }
}

/**
 * 这里为了简化状态管理，强制只有TERM状态的协程才可以重置，但其实刚创建好但没执行过的协程也应该允许重置的
 */
void Coroutine::reset(std::function<void()> cb) {
    LIBCO_ASSERT(m_stack);
    LIBCO_ASSERT(m_state == TERM);
    m_cb = cb;
    if (getcontext(&m_ctx)) {
        LIBCO_ASSERT2(false, "getcontext");
    }

    m_ctx.uc_link          = nullptr;
    m_ctx.uc_stack.ss_sp   = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;

    makecontext(&m_ctx, &Coroutine::MainFunc, 0);
    m_state = READY;
}

void Coroutine::resume() {
    LIBCO_ASSERT(m_state != TERM && m_state != RUNNING);
    SetThis(this);
    m_state = RUNNING;

    if (swapcontext(&(t_thread_coroutine->m_ctx), &m_ctx)) {
        LIBCO_ASSERT2(false, "swapcontext");
    }
}

void Coroutine::yield() {
    /// 协程运行完之后会自动yield一次，用于回到主协程，此时状态已为结束状态
    LIBCO_ASSERT(m_state == RUNNING || m_state == TERM);
    if (m_state != TERM) {
        m_state = READY;
    }
    LIBCO_LOG_DEBUG << "coroutine yield";
    
    SetThis(t_thread_coroutine.get());
    if (swapcontext(&m_ctx, &(t_thread_coroutine->m_ctx))) {
        LIBCO_ASSERT2(false, "swapcontext");
    }

}

/**
 * 这里没有处理协程函数出现异常的情况，同样是为了简化状态管理，并且个人认为协程的异常不应该由框架处理，应该由开发者自行处理
 */
void Coroutine::MainFunc() {
    Coroutine::ptr cur = GetThis(); // GetThis()的shared_from_this()方法让引用计数加1
    LIBCO_ASSERT(cur);

    cur->m_cb();
    LIBCO_LOG_DEBUG << "finish task";
    cur->m_cb    = nullptr;
    cur->m_state = TERM;

    auto raw_ptr = cur.get(); // 手动让t_coroutine的引用计数减1
    cur.reset();
    raw_ptr->yield();
}

} // namespace libco